/*
 * Copyright(C) 2021. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "Yolov8PostProcess.h"
#include "MxBase/Log/Log.h"
#include "MxBase/Maths/FastMath.h"
#include "MxBase/CV/ObjectDetection/Nms/Nms.h"



namespace {
    const int MODEL_INPUT_SIZE = 640;
    const int ALIGN_LEFT = 16;
    const int CONFIDENCE_IDX = 4;
    const int LABEL_START_OFFSET = 4;
    const int AVG_PARAM = 2;
}

Yolov8PostProcess &Yolov8PostProcess::operator=(const Yolov8PostProcess &other) {
    if (this == &other) {
        return *this;
    }
    ObjectPostProcessBase::operator=(other);
    objectnessThresh_ = other.objectnessThresh_; // Threshold of objectness value
    iouThresh_ = other.iouThresh_;
    paddingType_ = other.paddingType_;

    return *this;
}

APP_ERROR Yolov8PostProcess::Init(const std::map<std::string, std::shared_ptr<void>> &postConfig) 
{
    LogDebug << "Start to Init Yolov8PostProcess.";
    APP_ERROR ret = ObjectPostProcessBase::Init(postConfig);
    if (ret != APP_ERR_OK) {
        LogError << GetError(ret) << "Fail to superInit in ObjectPostProcessBase.";
        return ret;
    }

    LogDebug << "End to Init Yolov8PostProcess.";
    return APP_ERR_OK;
}

APP_ERROR Yolov8PostProcess::DeInit() {
    return APP_ERR_OK;
}

void Yolov8PostProcess::ConstructBoxFromOutput(float *output, size_t rows, std::vector<MxBase::ObjectInfo> &objectInfo,
    const MxBase::ResizedImageInfo &resizedImageInfo)
{
    double scale_ratio_w = (double)resizedImageInfo.widthOriginal / MODEL_INPUT_SIZE;
    double scale_ratio_h = (double)resizedImageInfo.heightOriginal / MODEL_INPUT_SIZE;

    for (size_t i = 0; i < rows; i++)
        {
            float maxProb = 0.f;
            int class_id;

            for (size_t j = 0; j < classNum_; j++)
            {
                float score = output[i + rows * (CONFIDENCE_IDX + j)];

                if (score > maxProb)
                {
                    maxProb = score;
                    class_id = j;
                }
            }
            
            if (maxProb > scoreThresh_)
            {
                float x = output[i + rows * 0];
                float y = output[i + rows * 1];
                float w = output[i + rows * 2];
                float h = output[i + rows * 3];
                
                float leftX = (x - w / AVG_PARAM) * scale_ratio_w;
                float leftY = (y - h / AVG_PARAM) * scale_ratio_h;
                float rightX = (x + w / AVG_PARAM) * scale_ratio_w;
                float rightY = (y + h / AVG_PARAM) * scale_ratio_h;

                MxBase::ObjectInfo obj;
                obj.x0 = leftX < 0.0 ? 0.0 : leftX;
                obj.y0 = leftY < 0.0 ? 0.0 : leftY;
                obj.x1 = rightX > resizedImageInfo.widthOriginal ? resizedImageInfo.widthOriginal : rightX;
                obj.y1 = rightY > resizedImageInfo.heightOriginal ? resizedImageInfo.heightOriginal : rightY;
                obj.confidence = maxProb;
                obj.classId = class_id;
                obj.className = configData_.GetClassName(obj.classId);
                if (maxProb < separateScoreThresh_[obj.classId])
                    return;
                objectInfo.push_back(obj);
            }
        }
}

// 将处理好的推理结果放入ObjectInfo
void Yolov8PostProcess::ObjectDetectionOutput(const std::vector<MxBase::TensorBase> &tensors,
                                              std::vector<std::vector<MxBase::ObjectInfo>> &objectInfos,
                                              const std::vector<MxBase::ResizedImageInfo> &resizedImageInfos) {
    LogDebug << "Yolov8PostProcess start to write results.";
    if (tensors.size() == 0) {
        return;
    }

    uint32_t batchSize = tensors[0].GetShape()[0];
    size_t rows = tensors[0].GetSize() / ((classNum_ + LABEL_START_OFFSET) * batchSize);
    for (size_t k = 0; k < batchSize; k++) {
        auto output = (float*)GetBuffer(tensors[0], k);
        std::vector<MxBase::ObjectInfo> objectInfo;
        ConstructBoxFromOutput(output, rows, objectInfo, resizedImageInfos[k]);
        MxBase::NmsSort(objectInfo, iouThresh_);
        objectInfos.push_back(objectInfo);
    }
    LogDebug << "Yolov8PostProcess write results successed.";
}

APP_ERROR Yolov8PostProcess::Process(const std::vector<MxBase::TensorBase> &tensors,
                                     std::vector<std::vector<MxBase::ObjectInfo>> &objectInfos,
                                     const std::vector<MxBase::ResizedImageInfo> &resizedImageInfos,
                                     const std::map<std::string, std::shared_ptr<void>> &paramMap) {
    LogDebug << "Start to Process Yolov8PostProcess.";
    APP_ERROR ret = APP_ERR_OK;
    if (resizedImageInfos.size() == 0) {
        ret = APP_ERR_INPUT_NOT_MATCH;
        LogError << GetError(ret) << "resizedImageInfos is not provided which is necessary for Yolov5PostProcess.";
        return ret;
    }
    auto inputs = tensors;
    ret = CheckAndMoveTensors(inputs);
    if (ret != APP_ERR_OK) {
        LogError << GetError(ret) << "CheckAndMoveTensors failed.";
        return ret;
    }
    
    ObjectDetectionOutput(inputs, objectInfos, resizedImageInfos);
    // LogObjectInfos(objectInfos);
    LogDebug << "End to Process Yolov8PostProcess.";
    return APP_ERR_OK;
}
